import taichi as ti

from ..material import Material
from ..utils import Vec3, Ray3, reflect, randUnitVec3


class Metal(Material):
    """类金属材质"""

    def __init__(self, fuzz: float = 0):
        super().__init__()
        Material.materials[self.index].attribute[0] = fuzz

    @ti.func
    def scatter(self: int, ray: Ray3, hitPoint: Vec3, norm: Vec3):
        ray.origin = hitPoint
        color = Material.getColor(self, hitPoint)
        if norm.dot(ray.direct) > 0:
            norm = -norm
        fuzz = Material.materials[self].attribute[0]
        ray.direct = reflect(ray.direct, norm)
        if fuzz != 0:
            ray.direct = ray.direct + fuzz * randUnitVec3()
            if ray.direct.dot(norm) <= 0:
                ray.direct = reflect(ray.direct, norm)
        return color, ray
